# +----------------------------------------------------------------------
# | ChatWork
# +----------------------------------------------------------------------
# | 这不是一个自由软件,您只能在不用于商业目的的前提下对程序代码进行修改和使用。
# | 任何企业和个人不允许对程序代码以任何形式任何目的再发布,商业使用请获取授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import json
import time
from typing import List
from tortoise.queryset import Q
from tortoise.transactions import in_transaction

from apps.api.service.agent.robot.teamed_service import RobotTeamedService
from common.service.subscribe_logic import SubscribeLogic
from exception import AppException
from common.enums.agent import AgentEnum
from common.utils.urls import UrlUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.models.users import UserModel
from common.models.agents import AgentMcpModel, AgentRobotModel
from common.models.agents import AgentMcpTeamedModel
from common.chain.mcp_server import MCPClient
from plugins.censor.driver import CensorDriver
from apps.api.schemas.agent.mcp import client_schema as schema
from apps.api.service.agent.mcp.teamed_service import McpTeamedService


class McpClientService:

    @classmethod
    async def whole(cls, user_id: int, params: schema.McpWholeIn) -> List[schema.McpWholeVo]:
        """
        所有MCP。

        Args:
            user_id (int): 用户ID。
            params (schema.McpWholeIn): 搜索参数。

        Returns:
            List[schema.McpWholeVo]: MCP列表Vo。

        Author:
            zero
        """
        uid: int = user_id

        # 智能体过滤
        if params.bot:
            # 查智能体
            robot = await AgentRobotModel.filter(code=params.bot).get()
            # 拥有者ID
            uid = robot.user_id
            # 验证权限
            await RobotTeamedService.has_permission(
                bot=str(params.bot),
                user_id=user_id,
                power=RobotTeamedService.POWER_VIEW
            )

        # 查询MCP
        lists = await (AgentMcpModel
                       .filter(user_id=uid)
                       .filter(is_delete=0)
                       .order_by("-sort", "-update_time")
                       .all())

        # 处理数据
        _lists = []
        for item in lists:
            _lists.append(schema.McpWholeVo(
                code=item.code,
                logo=await UrlUtil.to_absolute_url(item.logo),
                name=item.name,
                is_disable=item.is_disable,
            ))
        return _lists

    @classmethod
    async def get_tools(cls, url: str) -> List[dict]:
        """
        获取工具

        Args:
            url (str): 地址

        Returns:
            List[dict]

        Author:
            zero
        """
        client = MCPClient()
        try:
            await client.connect_to_server(service_url=url)
            return await client.get_tools_dict()
        except Exception as e:
            raise AppException(str(e))
        finally:
            await client.cleanup()

    @classmethod
    async def call_tool(cls, post: schema.McpCallIn):
        client = MCPClient()
        try:
            await client.connect_to_server(service_url=post.url)
            results = await client.call_tool(tool_name=post.tool_name, arguments=post.tool_args)
            content = []
            for properties in results.content:
                data = {}
                for k, v in properties:
                    if v is not None:
                        data[k] = v
                content.append(data)
            return {"content": content}
        except Exception as e:
            raise AppException(str(e))
        finally:
            await client.cleanup()

    @classmethod
    async def lists(cls, params: schema.McpSearchIn, user_id: int) -> List[schema.McpListVo]:
        """
        MCP列表

        Args:
            params (schema.McpSearchIn)
            user_id (int)

        Returns:
            List[schema.McpListVo]

        Author:
            zero
        """
        # 构建查询条件
        where = [Q(is_delete=0)]
        if params.keyword:
            where.append(Q(name__icontains=params.keyword))

        # 成员身份过滤
        permissions = {user_id: AgentEnum.POWER_OWNED}
        if params.owner == "private":
            where.append(Q(user_id=user_id))
        elif params.owner == "cooperate":
            team_lis = await AgentMcpTeamedModel.filter(user_id=user_id).all()
            team_ids = [item.mcp_id for item in team_lis if item.mcp_id]
            permissions = {team.user_id: team.power for team in team_lis}
            where.append(Q(id__in=team_ids or [int(0)]))
        else:
            team_lis = await AgentMcpTeamedModel.filter(user_id=user_id).all()
            team_ids = [int(item.mcp_id) for item in team_lis if item.mcp_id]
            if team_ids:
                permissions = {team.user_id: team.power for team in team_lis}
                where.append(Q(user_id=user_id) | Q(id__in=team_ids or [int(0)]))
            else:
                where.append(Q(user_id=user_id))

        # 查询MCP列表
        _model = AgentMcpModel.filter(*where).order_by("-sort", "-update_time")
        _pager = await AgentMcpModel.paginate(
            model=_model,
            page_no=params.page_no,
            page_size=params.page_size,
            fields=["id", "code", "user_id", "logo", "name", "intro", "is_disable", "create_time", "update_time"],
            datetime_format="%m/%d %H:%M"
        )

        # 查询用户信息
        user_ids = [item["user_id"] for item in _pager.lists if item["user_id"]]
        users = await UserModel.queryUserDictsById(user_ids)

        _lists = []
        for item in _pager.lists:
            logo: str = await UrlUtil.to_absolute_url(item["logo"])
            owner: bool = True if item["user_id"] == user_id else False
            power: int = permissions.get(item["user_id"], 0)
            user: dict = users.get(item["user_id"], {})

            _lists.append(schema.McpListVo(
                code=item["code"],
                logo=logo,
                name=item["name"],
                intro=item["intro"],
                power=power,
                is_owner=owner,
                is_disable=item["is_disable"],
                create_time=item["create_time"],
                update_time=item["update_time"],
                user={
                    "sn": user.get("sn", ""),
                    "nickname": user.get("nickname", ""),
                    "avatar": await UrlUtil.to_absolute_url(user.get('avatar', ""))
                }
            ))

        return _lists

    @classmethod
    async def detail(cls, code: str, user_id: int) -> schema.McpDetailVo:
        """
        MCP详情

        Args:
            code (str): 编号
            user_id (int): 用户ID

        Returns:
            schema.McpDetailVo

        Author:
            zero
        """
        # 查询MCP
        mcp = await AgentMcpModel.filter(code=code, is_delete=0).get()

        # 权限的验证
        if mcp.user_id != user_id:
            await McpTeamedService.has_permission(mcp.code, user_id, AgentEnum.POWER_VIEW)

        # 权限查询
        power: int = McpTeamedService.POWER_OWNED
        if mcp.user_id != user_id:
            power = await McpTeamedService.has_permission(mcp.code, user_id, McpTeamedService.POWER_VIEW)

        # 返回结果
        return schema.McpDetailVo(
            logo=await UrlUtil.to_absolute_url(mcp.logo),
            code=mcp.code,
            name=mcp.name,
            intro=mcp.intro,
            url=mcp.url,
            power=power,
            is_owner=True if mcp.user_id == user_id else False,
            create_time=TimeUtil.timestamp_to_date(mcp.create_time),
            update_time=TimeUtil.timestamp_to_date(mcp.update_time),
            tools=json.loads(mcp.tools or "[]"),
        )

    @classmethod
    async def add(cls, post: schema.McpAddIn, user_id: int) -> None:
        """
        MCP新增

        Args:
            post (schema.McpAddIn): 新增参数
            user_id (int): 用户ID

        Author:
            zero
        """
        # 验证剩余数
        surplus: int = await SubscribeLogic.get_surplus_quota(user_id, scene="mcp")
        if surplus <= 0:
            raise AppException("MCP创建额度不足!")

        # 违规审查
        review = CensorDriver.REVIEW_TEXTS
        censor = await CensorDriver.text(user_id=int(user_id), text=post.name, scene=review)
        if censor is not None:
            raise AppException(f"MCP名称违规: {censor}")

        # 创建数据
        logo: str = "static/images/default_mcp.png"
        code: str = await ToolsUtil.make_rand_sn(AgentMcpModel, field="code", length=16, scene="char")
        await AgentMcpModel.create(
            user_id=user_id,
            logo=logo,
            code=code,
            name=post.name,
            url=post.url,
            tools=json.dumps(post.tools, ensure_ascii=False),
            create_time=int(time.time()),
            update_time=int(time.time())
        )

    @classmethod
    async def edit(cls, post: schema.McpEditIn, user_id: int) -> None:
        """
        MCP编辑

        Args:
            post (schema.McpAddIn): 编辑参数
            user_id (int): 用户ID

        Author:
            zero
        """
        # 违规审查
        review = CensorDriver.REVIEW_TEXTS
        values = post.name + "\n" + post.intro
        censor_info = await CensorDriver.text(user_id=int(user_id), text=values, scene=review)
        if censor_info is not None:
            raise AppException(f"MCP信息违规: {str(censor_info)}")

        # 编辑数据
        mcp = await AgentMcpModel.filter(code=post.code, is_delete=0).get()
        await AgentMcpModel.filter(id=mcp.id).update(
            logo=await UrlUtil.to_absolute_url(post.logo),
            name=post.name,
            intro=post.intro,
            url=post.url,
            tools=json.dumps(post.tools, ensure_ascii=False),
            update_time=int(time.time())
        )

    @classmethod
    async def delete(cls, code: str, user_id: int) -> None:
        """
        MCP删除

        Args:
            code (schema.McpAddIn): MCP编号
            user_id (int): 用户ID

        Author:
            zero
        """
        # 数据验证
        mcp = await AgentMcpModel.filter(code=code, is_delete=0).get()

        # 权限的校验
        await McpTeamedService.has_permission(mcp.code, user_id, McpTeamedService.POWER_OWNED)

        async with in_transaction("mysql"):
            await AgentMcpModel.filter(id=mcp.id).update(
                is_delete=1,
                delete_time=int(time.time())
            )
